﻿using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Security.Cryptography.X509Certificates;
using System.Text;
using System.Threading.Tasks;
using RayDen.Library.Components.Surface;
using RayDen.Library.Core;
using RayDen.Library.Core.Primitives;
using RayDen.RayEngine.Core.Interface;
using RayDen.RayEngine.Core.Types;
using RayDen.RayEngine.Engines.PathTracer;
using RayDen.RayEngine.Samplers;

namespace RayDen.RayEngine.Engines.PathSpaceMLT
{
    [Flags]
    public enum MltPathSamplerState
    {
        Init,
        Mutated
    }


    public class MetropolisPathSampler : PathSamplerBase
    {

        public int MaxMutations = 100;

        internal RgbSpectrum Throughput, XRadiance;
        public PathTracerPathState PathState;

        protected internal bool mutate;
        protected int mutationCount;
        protected int depth, tracedShadowRayCount;
        protected BsdfEvent bsdfEvent, prevEvent, shadowRayEvent;
        protected float traveledDistance, mutationStrategyWeight;
        protected float pathDensity, xPathDensity, x_b, y_b;
        public UInt64 PathExpression;

        protected internal RayEngineScene scene;
        protected internal SurfaceIntersectionData hitInfo = null;
        protected internal PathVertex[] Vertices, Xvertices;
        protected internal Point pathStart;
        public int MaxRaysPerPath
        {
            get;
            set;
        }


        public override void InitPath(IPathProcessor buffer)
        {
            base.InitPath(buffer);
            this.scene = pathIntegrator.Scene;

            if (Vertices == null)
            {
                Vertices = new PathVertex[scene.MaxPathDepth * 3];
                Xvertices = new PathVertex[scene.MaxPathDepth * 3];
            }
            this.pathDensity = 1.0f;
            this.Radiance = new RgbSpectrum(0f);
            this.Throughput = new RgbSpectrum(1f);
            this.PathState = PathTracerPathState.EyeVertex;
            if (!mutate)
            {
                this.Sample = pathIntegrator.Sampler.GetSample(this.Sample);
                pathIntegrator.Scene.Camera.GetRay(Sample.imageX, Sample.imageY, out this.PathRay);
                pathStart = PathRay.Org;
                XRadiance = new RgbSpectrum();
                xPathDensity = 1.0f;
            }
            else
            {
                PathRay.Org = pathStart;
                PathRay.Dir = -Vertices[0].Wo;
                this.BackupPath();
                this.MutatePath();
                mutationCount++;
                if (mutationCount > MaxMutations)
                {
                    this.mutate = false;
                    this.mutationCount = 0;
                    this.InitPath(buffer);
                }
            }
            this.RayIndex = -1;
            this.depth = 0;
            this.bsdfEvent = BsdfEvent.Specular;
            this.prevEvent = BsdfEvent.None;
            shadowRayEvent = BsdfEvent.None;
        }

        public override void Splat(SampleBuffer sampleBuffer)
        {
            if (!mutate)
            {
                if (!Radiance.IsBlack())
                {
                    XRadiance = Radiance;
                    xPathDensity = pathDensity;
                }
                sampleBuffer.SplatSample(this.Sample.imageX, this.Sample.imageY, ref this.Radiance);
                this.InitPath(pathIntegrator);
            }
            else
            {
                if (!Radiance.IsBlack())
                {

                    var Tyx = pathDensity;
                    var Txy = xPathDensity;
                    var FY = Radiance.y();
                    var FX = XRadiance.y();
                    Radiance /= FY;
                    var Axy = Math.Min(1, (FY*Txy)/(FX*Tyx));
                    if (Sample.GetLazyValue() < Axy)
                    {
                        xPathDensity = pathDensity;
                        XRadiance = Radiance;
                    }
                    else
                    {
                        RestorePath();
                    }
                }
                else
                {
                    RestorePath();
                }
                var rad = this.XRadiance * (1f / MaxMutations);
                sampleBuffer.SplatSample(this.Sample.imageX, this.Sample.imageY, ref rad);
                this.InitPath(pathIntegrator);
            }


        }


        private void MutatePath()
        {
            pixelOffset(Sample.GetLazyValue(), Sample.GetLazyValue(), ref Sample.imageX, ref Sample.imageY);
            for (int index = 0; index < Vertices.Length; index++)
            {
                var pathVertex = Vertices[index];
                if (pathVertex == null)
                    continue;
                if (pathVertex.Bsdf!=null && pathVertex.Bsdf.IsDiffuse())
                {
                    angularOffset(0.0001f, 0.1f, ref pathVertex.Wi);
                }
            }
            mutationStrategyWeight = 1.0f;

        }
        void pixelOffset(float r1, float r2, ref float x, ref float y)
        {
            x_b = x;
            y_b = y;
            float phi = Sample.GetLazyValue() * 2f * MathLab.M_PI;
            float r = r2 * MathLab.Exp(-(float)Math.Log(r2 / r1) * Sample.GetLazyValue());
            x += r * MathLab.Cos(phi);
            y += r * MathLab.Sin(phi);
        }

        void pixelOffset(float r1, float r2, ref double x, ref double y)
        {
            x_b = (float)x;
            y_b = (float)y;
            float phi = Sample.GetLazyValue() * 2f * MathLab.M_PI;
            float r = r2 * MathLab.Exp(-(float)Math.Log(r2 / r1) * Sample.GetLazyValue());
            x += r * MathLab.Cos(phi);
            y += r * MathLab.Sin(phi);
        }
        void angularOffset(float theta1, float theta2, ref Vector N)
        {
            // Make a UVN coordinate system from N
            var vx = new Vector(1, 0, 0);
            var vy = new Vector(0, 1, 0);
            Vector U, V;


            if (Math.Abs(N.x) < 0.5f) U = Vector.Cross(ref N, ref vx);
            else
                U = Vector.Cross(ref N, ref vy);
            U.Normalize();
            V = Vector.Cross(ref U, ref N);
            // Determine offsets using the approximation θ ≈ sinθ
            float phi = Sample.GetLazyValue() * 2f * MathLab.M_PI;

            float r = theta2 * MathLab.Exp(-(float)Math.Log(theta2 / theta1) * Sample.GetLazyValue()); ;
            // Calculate the new direction
            N = N + r * MathLab.Cos(phi) * U + r * MathLab.Sin(phi) * V;
            N.Normalize();
        }

        private void BackupPath()
        {
            for (int index = 0; index < Vertices.Length; index++)
            {
                var pathVertex = Vertices[index];
                if (pathVertex != null)
                {
                    Xvertices[index] = (PathVertex)pathVertex.Clone();
                }
            }
        }

        private void RestorePath()
        {
            Sample.imageX = x_b;
            Sample.imageY = y_b;
            for (int index = 0; index < Xvertices.Length; index++)
            {
                var pathVertex = Xvertices[index];
                if (pathVertex != null)
                {
                    Vertices[index].HitPoint = pathVertex.HitPoint;
                    Vertices[index].Bsdf = pathVertex.Bsdf;
                    Vertices[index].BsdfSample = pathVertex.BsdfSample;
                    Vertices[index].BsdfWeight = pathVertex.BsdfWeight;
                    Vertices[index].Event = pathVertex.Event;
                    Vertices[index].ShadeNormal = pathVertex.ShadeNormal;
                    Vertices[index].ThroughtPut = pathVertex.ThroughtPut;
                    Vertices[index].TriangleIndex = pathVertex.TriangleIndex;
                    Vertices[index].Wi = pathVertex.Wi;
                    Vertices[index].Wo = pathVertex.Wo;
                }
            }
        }

        public override bool FillRayBuffer(RayBuffer rayBuffer)
        {
            var leftSpace = rayBuffer.LeftSpace();
            if (((PathState == PathTracerPathState.EyeVertex) && (1 > leftSpace)) ||
            ((PathState == PathTracerPathState.ShadowRaysOnly) && (tracedShadowRayCount > leftSpace)) ||
            ((PathState == PathTracerPathState.NextVertex) && (tracedShadowRayCount + 1 > leftSpace)))
                return false;
            if (PathState != PathTracerPathState.ShadowRaysOnly)
                RayIndex = rayBuffer.AddRay(ref PathRay);
            return true;
        }



        public override void Advance(RayBuffer rayBuffer, SampleBuffer consumer)
        {
            int currentTriangleIndex = 0;
#if VERBOSE

            try {
#endif

            if (!mutate)
                AdvanceNewPath(rayBuffer, consumer);
            else
            {
                AdvanceMutatedPath(rayBuffer, consumer);
            }


#if VERBOSE
            }
            catch (Exception ex) {
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine("Advance path exception");
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine("Error triangle {0}", currentTriangleIndex);
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine(ex.Message);
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine(ex.StackTrace);

            }
#endif
        }

        private void AdvanceMutatedPath(RayBuffer rayBuffer, SampleBuffer consumer)
        {
            int currentTriangleIndex;
            RayHit rayHit = rayBuffer.rayHits[RayIndex];
            Vector wo = -PathRay.Dir;
            depth++;
            bool missed = rayHit.Index == 0xffffffffu;
            if (missed || PathState == PathTracerPathState.ShadowRaysOnly)
            {
                if (missed)
                {
                    //Radiance += this.scene.SampleEnvironment(ref wo) * Throughput;
                    var sampledEnvironment = this.scene.SampleEnvironment(ref wo);
                    Radiance.MAdd(ref sampledEnvironment, ref Throughput);
                }
                Splat(consumer);

                return;
            }

            var idx = Vertices.Count(vv => vv != null)-1;
            //depth > scene.MaxPathDepth || 
            if (depth == (idx))
            {

                var lt = scene.GetLightByIndex(Vertices[idx].TriangleIndex);
                if (lt != null)
                {
                    float pdf;
                    var le = hitInfo.Color * lt.Emittance(ref wo, out pdf);
                    //Radiance += Throughput * le;

                    Radiance.MAdd(ref Throughput, ref le);
                }

                Splat(consumer);

                return;
            }

            //var mesh = scene.GetMeshByTriangleIndex(currentTriangleIndex);
            //rayHit.Index += (uint)mesh.StartTriangle;
            // Something was hit
            var bsdf = EvalIntersection(ref rayHit);
            currentTriangleIndex = (int)rayHit.Index;

            if (hitInfo == null)
            {
                Debugger.Break();
            }
            //If Hit light)

            if (hitInfo.IsLight)
            {
                //if (bsdfEvent.Has(BsdfEvent.Specular) || depth == 1)
                {
                    var lt = scene.GetLightByIndex(currentTriangleIndex);
                    if (lt != null)
                    {
                        float pdf;
                        var le = hitInfo.Color * lt.Emittance(ref wo, out pdf);
                        //Radiance += Throughput * le;

                        Radiance.MAdd(ref Throughput, ref le);
                    }
                }
                Splat(consumer);

                return;
            }

            var hitPoint = PathRay.Point(rayHit.Distance);
            tracedShadowRayCount = 0;


            if (!hitInfo.TextureData.Alpha.IsBlack())
            {
                Throughput *= (RgbSpectrum.UnitSpectrum() - (RgbSpectrum)hitInfo.TextureData.Alpha);
                PathRay = new RayData(hitPoint, -wo);
                return;
            }

            float fPdf;
            Vector wi;
            RgbSpectrum f;
            hitInfo.TextureData.Throughput = Throughput;
            prevEvent = bsdfEvent;


            f = bsdf.Sample_f(ref wo, out wi, ref hitInfo.Normal, ref hitInfo.ShadingNormal, ref Throughput,
                Sample.GetLazyValue(), Sample.GetLazyValue(), Sample.GetLazyValue()
                , ref hitInfo.TextureData, out fPdf, out bsdfEvent);

            //if (prevEvent.Has(BsdfEvent.Transmit) && bsdfEvent.Has(BsdfEvent.Transmit))
            //{
            //    Throughput *= MathLab.Exp(-rayHit.Distance) * 0.2f;
            //    pathDensity *= MathLab.Exp(-rayHit.Distance) * 0.2f;
            //}

            
            if (Vertices[depth-1].Event != bsdfEvent || bsdf != Vertices[depth-1].Bsdf)
            {
                Splat(consumer);
                return;
            }
            if (bsdfEvent.Has(BsdfEvent.Specular))
                wi = Vertices[depth-1].Wi;

            if ((fPdf <= 0.0f) || f.IsBlack())
            {
                if (tracedShadowRayCount > 0)
                    PathState = PathTracerPathState.ShadowRaysOnly;
                else
                {
                    Splat(consumer);
                }
                return;
            }
            Throughput *= (f * hitInfo.Color) / fPdf;
            pathDensity *= fPdf;
            if (depth > scene.MaxPathDepth)
            {
                float prob = Math.Max(Throughput.Filter(), scene.RussianRuletteImportanceCap);
                if (prob >= Sample.GetLazyValue())
                {
                    Throughput /= prob;
                    pathDensity *= prob;
                }
                else
                {
                    if (tracedShadowRayCount > 0)
                        PathState = PathTracerPathState.ShadowRaysOnly;
                    else
                    {
                        Splat(consumer);
                    }

                    return;
                }
            }

            PathRay.Org = hitPoint;
            PathRay.Dir = wi.Normalize();
            PathState = PathTracerPathState.NextVertex;
        }

        private void AdvanceNewPath(RayBuffer rayBuffer, SampleBuffer consumer)
        {
            int currentTriangleIndex;
            RayHit rayHit = rayBuffer.rayHits[RayIndex];
            Vector wo = -PathRay.Dir;
            PathVertex vertex = null;
            if (Vertices[depth] == null)
            {
                vertex = new PathVertex();
                Vertices[depth] = vertex;
            }
            else
            {
                vertex = Vertices[depth];
            }
            depth++;
            bool missed = rayHit.Index == 0xffffffffu;
            if (missed || PathState == PathTracerPathState.ShadowRaysOnly || depth > scene.MaxPathDepth)
            {
                if (missed)
                {
                    //Radiance += this.scene.SampleEnvironment(ref wo) * Throughput;
                    var sampledEnvironment = this.scene.SampleEnvironment(ref wo);
                    Radiance.MAdd(ref sampledEnvironment, ref Throughput);
                    if (scene.EnvironmentMap != null)
                    {
                        mutate = true;
                    }
                }
                Splat(consumer);

                return;
            }

            //var mesh = scene.GetMeshByTriangleIndex(currentTriangleIndex);
            //rayHit.Index += (uint)mesh.StartTriangle;
            // Something was hit
            var bsdf = EvalIntersection(ref rayHit);
            currentTriangleIndex = (int)rayHit.Index;
            vertex.Bsdf = bsdf;
            vertex.TriangleIndex = currentTriangleIndex;

            if (hitInfo == null)
            {
                Debugger.Break();
            }
            //If Hit light)

            if (hitInfo.IsLight)
            {
                //if (bsdfEvent.Has(BsdfEvent.Specular) || depth == 1)
                {
                    var lt = scene.GetLightByIndex(currentTriangleIndex);
                    if (lt != null)
                    {
                        float pdf;
                        var le = hitInfo.Color * lt.Emittance(ref wo, out pdf);
                        //Radiance += Throughput * le;
                        if (!le.IsBlack())
                            mutate = true;

                        Radiance.MAdd(ref Throughput, ref le);
                    }
                }
                Splat(consumer);

                return;
            }

            var hitPoint = PathRay.Point(rayHit.Distance);
            vertex.HitPoint = hitPoint;
            tracedShadowRayCount = 0;


            if (!hitInfo.TextureData.Alpha.IsBlack())
            {
                Throughput *= (RgbSpectrum.UnitSpectrum() - (RgbSpectrum)hitInfo.TextureData.Alpha);
                PathRay = new RayData(hitPoint, -wo);
                return;
            }

            float fPdf;
            Vector wi;
            RgbSpectrum f;
            hitInfo.TextureData.Throughput = Throughput;
            prevEvent = bsdfEvent;


            f = bsdf.Sample_f(ref wo, out wi, ref hitInfo.Normal, ref hitInfo.ShadingNormal, ref Throughput,
                Sample.GetLazyValue(), Sample.GetLazyValue(), Sample.GetLazyValue()
                , ref hitInfo.TextureData, out fPdf, out bsdfEvent);

            //if (prevEvent.Has(BsdfEvent.Transmit) && bsdfEvent.Has(BsdfEvent.Transmit))
            //{
            //    Throughput *= MathLab.Exp(-rayHit.Distance) * 0.2f;
            //}
            vertex.BsdfSample = f;
            vertex.BsdfWeight = fPdf;
            vertex.Event = bsdfEvent;
            vertex.Wo = wo;
            vertex.Wi = wi;
            vertex.ShadeNormal = hitInfo.ShadingNormal;


            if ((fPdf <= 0.0f) || f.IsBlack())
            {
                if (tracedShadowRayCount > 0)
                    PathState = PathTracerPathState.ShadowRaysOnly;
                else
                {
                    Splat(consumer);
                }
                return;
            }
            Throughput *= (f * hitInfo.Color) / fPdf;
            pathDensity *= fPdf;

            if (depth > scene.MaxPathDepth)
            {
                float prob = Math.Max(Throughput.Filter(), scene.RussianRuletteImportanceCap);
                if (prob >= Sample.GetLazyValue())
                {
                    Throughput /= prob;
                    pathDensity *= prob;

                }
                else
                {
                    if (tracedShadowRayCount > 0)
                        PathState = PathTracerPathState.ShadowRaysOnly;
                    else
                    {
                        Splat(consumer);
                    }

                    return;
                }
            }

            PathRay.Org = hitPoint;
            PathRay.Dir = wi.Normalize();
            PathState = PathTracerPathState.NextVertex;
        }

        private ISurfaceMaterial EvalIntersection(ref RayHit rayHit)
        {
            if (hitInfo == null)
            {
                hitInfo = SurfaceSampler.GetIntersection(ref PathRay, ref rayHit);
            }
            else
            {
                SurfaceSampler.GetIntersection(ref PathRay, ref rayHit, ref hitInfo);
            }

            return hitInfo.MMaterial;
        }

        protected override bool ShadowRayTest(ref RayHit shadowRayHit, ref RayData shadowRay, out RgbSpectrum attenuation, out bool continueTrace)
        {
            var hit = shadowRayHit.Index == 0xffffffffu || scene.IsLight((int)shadowRayHit.Index);
            attenuation = new RgbSpectrum(1f);
            continueTrace = false;

            if (hit) return true;
            var mesh = scene.GetMeshByTriangleIndex((int)shadowRayHit.Index);
            var mat =
                //SurfaceMaterials.CreateMaterial(scene.MaterialProvider.Get(mesh.MaterialName));
                scene.MatLib.GetSurfMat(mesh != null ? mesh.MaterialName : "default");
            shadowRayEvent = mat.Type.TypeToEvent();
            if (mat.TransparentShadows)
            {
                continueTrace = true;
                return true;
            }

            if (shadowRayEvent.Has(BsdfEvent.Transmit))
            {
                var att = MathLab.Exp(-0.1f * shadowRayHit.Distance);
                attenuation *= att + att * MathLab.Exp(0.1f) * hitInfo.TextureData.Transmittance;
                continueTrace = true;
                return true;
            }

            if (mat.AlphaTexture != null)
            {
                continueTrace = true;
                this.SurfaceSampler.GetIntersection(ref shadowRay, ref shadowRayHit, ref hitInfo, IntersectionOptions.ResolveTextures);
                attenuation = (RgbSpectrum)hitInfo.TextureData.Alpha;
                hit = true;
                //mat.MaterialData.Kt.Filter() < Sample.GetLazyValue();
            }

            return hit;
        }

    }
}
